import requests
import base64
import os
import sys
from PIL import Image
from io import BytesIO
from main import VIDEO_API_KEY,VIDEO_API_URL,VIDEO_SAVE_PATH,VIDEO_MODEL_NAME


class vid_model:
    def __init__(self, api_key, end_point, vid_model_name, vid_save_path):
        self.VID_API_KEY = api_key
        self.END_POINT = end_point
        self.VID_MODEL_NAME = vid_model_name
        self.PATH = vid_save_path

    def save_vid(self, url, save_path):
        try:
            response = requests.get(url, stream=True)
            if response.status_code == 200:
                # 自动创建目录
                os.makedirs(os.path.dirname(save_path), exist_ok=True)

                # 写入文件
                with open(save_path, 'wb') as f:
                    for chunk in response.iter_content(1024):
                        f.write(chunk)
                print(f"视频已保存至：{save_path}")
            else:
                print(f"下载失败，状态码：{response.status_code}")
        except Exception as e:
            print(f"发生错误：{str(e)}")

    def vid_gen(self, prompt, img, is_tail=0):
        headers = {
            "Authorization": f"Bearer {self.VID_API_KEY}",
            "Content-Type": "application/json"
        }
        if is_tail:
            payload = {
                "model_name": self.VID_MODEL_NAME,
                "prompt": prompt,
                "image": base64.b64encode(img).decode("utf-8"),
            }
        else:
            payload = {
                "model_name": self.VID_MODEL_NAME,
                "prompt": prompt,
                "image_tail": base64.b64encode(img).decode("utf-8"),
            }
        response = requests.post(self.END_POINT, headers=headers, json=payload)
        if response.status_code == 200:
            image_url = response.json()['data']['task_result']['videos']['url']
            self.save_vid(image_url, self.PATH)


vid_client = vid_model(VIDEO_API_KEY, VIDEO_API_URL, VIDEO_MODEL_NAME, VIDEO_SAVE_PATH)

if len(sys.argv) == 4 or len(sys.argv) == 3:
    if len(sys.argv) == 3:
        Prompt = sys.argv[1]
        Is_tail = sys.argv[2]
        buffer = None
    else:
        Prompt = sys.argv[1]
        Is_tail = sys.argv[2]
        buffer = BytesIO()
        Img = Image.open(sys.argv[3])
        Img.save(buffer, format="JPEG")
        buffer.seek(0)
    vid_client.vid_gen(prompt=Prompt, is_tail=bool(Is_tail), img=buffer)
else:
    print("false")